﻿namespace FSViterbi
open System
open Microsoft.FSharp.Collections
type dict<'a, 'b> = System.Collections.Generic.Dictionary<'a,'b>

type Viterbi<'s,'o when 's:equality>
    (startingProbs:     's -> float,
     transitionProbs :  's * 's -> float,
     emissionProbs:     's * 'o -> float) =

    let mutable lastV = dict<'s, float>()
    let mutable targetStates : 's seq = Seq.empty
    let mutable sourceStates : 's seq = Seq.empty
    let mutable _v = [dict<'s, float>()]
    let path = dict<'s,'s list>()

    let initialize (obs,states) =
        targetStates <- states

        for state in states do
            lastV.[state] <- (startingProbs state) * emissionProbs(state,obs)
            path.[state] <- [state]

        _v <- [lastV]

    let nextObservation (obs,states) =
        sourceStates <- targetStates
        targetStates <- states
        let nextV = dict<'s, float>()

        //parallelized implementation
        //replace 'PSeq' with 'Seq' for serial version
        targetStates |> PSeq.iter (fun target ->
            let arg,prob = 
                sourceStates 
                |> PSeq.map (fun source -> 
                    let a = lastV.[source]
                    let b = transitionProbs(source,target)
                    let c = emissionProbs(target,obs)
                    source,a*b*c)
                |> PSeq.maxBy (fun (_,prob) -> prob)
            nextV.[target] <- prob
            path.[target] <- target::path.[arg]) //copying not needed

        lastV <- nextV
        _v <- nextV::_v

    member x.V = _v

    interface IObserver<'o*'s seq> with 
       member x.OnError e = raise e
       member x.OnCompleted() = ()
       member x.OnNext value = 
        if Seq.isEmpty targetStates then 
            initialize value
        else
            nextObservation value
